# -*- coding: utf-8 -*-
# @Author   : zxc19
# @Time     : 2023/7/6 19:54
# @FileName : yaml_utils.py
# @Desc     : 请输入功能模块描述
import yaml
from loguru import logger
from typing import Any, Dict, Union, List
from pathlib import Path


class YamlUtil(object):
    def __init__(self, file_path: Union[Path, str]) -> None:
        """
        初始化YAML工具类
        :param file_path: yaml文件路径
        """
        self.file_path = Path(file_path)

    def read_yaml(self):
        """
        读取yaml文件数据
        :return:
        """
        try:
            logger.debug(f"读取YAML文件: {self.file_path}")
            with open(self.file_path, 'r', encoding='utf-8') as file:
                yaml_data = yaml.safe_load(file)
            logger.info(f"成功读取YAML文件: {self.file_path}")
            return yaml_data
        except FileNotFoundError:
            logger.error(f"错误: 文件 '{self.file_path}' 不存在.")
        except yaml.YAMLError as e:
            logger.error(f"错误: 解析YAML文件 '{self.file_path}' 失败: {e}")
        except Exception as e:
            logger.error(f"错误: 解析YAML文件 '{self.file_path}' 失败: {e}")
        return None

    def write_yaml(self, data: Dict[str, Any] | List) -> None:
        """
        写入数据到yaml文件
        :param data: 待写入数据
        :return:
        """
        try:
            with open(self.file_path, 'w', encoding='utf-8' ) as file:
                yaml.dump(data, file, indent=2, default_flow_style=False, encoding='utf-8', allow_unicode=True,
                          sort_keys=False)
            logger.info(f"成功写入YAML文件: {self.file_path}")
        except yaml.YAMLError as e:
            logger.error(f"错误: 写入YAML文件 '{self.file_path}' 失败: {e}")
        except Exception as e:
            logger.error(f"错误: 写入YAML文件 '{self.file_path}' 失败: {e}")


if __name__ == '__main__':
    data = {
        "Windows主机": ["登录成功审计", "登录失败审计", "用户注销审计", "帐户变更审计", "口令变更审计",
                        "帐户权限变更审计", "策略修改审计", "系统文件删除审计", "可执行文件安装审计", "启动进程审计",
                        "系统重启审计", "拒绝可疑连接行为审计", "网络蠕虫行为审计"],
        "Linux/Unix主机": ["登录成功审计", "登录失败审计", "用户注销审计", "帐户变更审计", "口令变更审计",
                           "帐户权限变更审计", "策略修改审计", "系统文件删除审计", "可执行文件安装审计", "启动进程审计",
                           "系统重启审计", "SU成功动作审计", "SU失败动作审计", "su root动作审计"],
        "防火墙": ["登录成功审计", "拒绝可疑连接行为审计"],
        "扫描器": ["网络探测威胁行为审计"],
        "IDS/IPS": ["网络蠕虫行为审计", "网络拒绝服务行为审计"],
        "防病毒": ["普通病毒感染告警审计", "高风险病毒感染告警审计"],
        "数据库": ["数据库高危操作审计", "数据库Insert动作审计", "数据库Delete动作审计", "数据库Create动作审计",
                   "数据库Drop动作审计", "数据库Update动作审计", "数据库Select动作审计", "数据库用户登录审计",
                   "数据库用户帐号变更审计"],
        "萨班斯": ["登录成功审计", "登录失败审计", "用户注销审计", "数据库用户登录审计", "帐户变更审计",
                   "审计策略变更审计"],
        "等级保护": ["审计策略变更审计", "主机访问控制-Linux", "主机访问控制-Windows", "应用访问控制", "数据库访问控制",
                     "网络设备访问控制", "网络连接分布", "设备故障分布", "防火墙访问控制"]
    }
    num = 1
    res = []
    for k, v in data.items():
        for d in v:
            res.append({
                "case_id": f"add_by_template-{str(num).zfill(3)}",
                'module': '审计策略',
                'function': '内置模板添加',
                'level': 0,
                'is_run': 'Y',
                'case_title': f'内置模板添加审计策略-{k}-{d}',
                'template_type': k,
                'template_name': d,
                'expects': {
                    'message': [
                        {
                            'check_point': 'content',
                            'check_type': 'eq',
                            'check_value': '操作成功'
                        }
                    ]
                }
            })
            num += 1
    print(res)

    YamlUtil('../test_data/strategy_manager/audit_strategy/add_by_template.yaml').write_yaml(res)
    print(YamlUtil('../test_data/strategy_manager/audit_strategy/add_by_template.yaml').read_yaml())
